Add vectorization in elementwise_util #9432
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
swolchok
requested review from
JacobSzwejbka,
lucylq and
manuelcandales
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/9432
Note: Links to docs will display an error until the docs builds have been completed. ❌ 70 New FailuresAs of commit 0beabbb with merge base 1572381 ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This was referenced Mar 20, 2025
Merged
swolchok
added a commit
that referenced
this pull request
Mar 20, 2025
this works with op_mul, which is vectorized-friendly, but doesn't work when we roll out to pattern.h because those ops will not work with Vectorized yet. See TODO in elementwise_util.h ghstack-source-id: 30d2311bed080c3a5390ab00ca20a1e33563f077 ghstack-comment-id: 2738665976 Pull Request resolved: #9432
swolchok
added a commit
that referenced
this pull request
Apr 2, 2025
this works with op_mul, which is vectorized-friendly, but doesn't work when we roll out to pattern.h because those ops will not work with Vectorized yet. See TODO in elementwise_util.h ghstack-source-id: 8d76653f819dc58a0c93540f3d71a89bfdb7cd26 ghstack-comment-id: 2738665976 Pull Request resolved: #9432
|
This PR is not yet complete (I need to go through and make sure no ops are needlessly held back from vectorization, and I need to overload unary - on at::vec::Vectorized so that op_sigmoid.cpp can vectorize nicely), but it builds and should vectorize a large chunk of elementwise portable ops. Hopefully this makes the direction/vision more clear @kimishpatel (CI failures are expected; the builds are contingent on pytorch/pytorch#150380 which is not yet actually merged into PyTorch core. We'll need a pin bump.) |
swolchok
commented
Apr 2, 2025
kernels/portable/cpu/op_where.cpp
Outdated
Comment on lines
50
to
52
| [](const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b, const CTYPE_COMPUTE val_c) { | ||
| return val_c ? val_a : val_b; | ||
| }, |
There was a problem hiding this comment.
(ATen, and therefore our own optimized op_where, doesn't vectorize this)
I am not fully sold but I havent looked at the PR in detail so I will further comment after that. Maybe you are right in that not everything will fit the Vectorized pattern and for those cases your pattern of checking callable for vectorized is probably a better way to do it. |
This was referenced Apr 2, 2025
swolchok
added a commit
that referenced
this pull request
Apr 2, 2025
this works with op_mul, which is vectorized-friendly, but doesn't work when we roll out to pattern.h because those ops will not work with Vectorized yet. See TODO in elementwise_util.h ghstack-source-id: 8d76653f819dc58a0c93540f3d71a89bfdb7cd26 ghstack-comment-id: 2738665976 Pull Request resolved: #9432
This was referenced Apr 2, 2025
swolchok
added a commit
that referenced
this pull request
Apr 2, 2025
this works with op_mul, which is vectorized-friendly, but doesn't work when we roll out to pattern.h because those ops will not work with Vectorized yet. See TODO in elementwise_util.h ghstack-source-id: 033b63ce3bee8a0136efdab3e03905cafb79b915 ghstack-comment-id: 2738665976 Pull Request resolved: #9432
swolchok
added a commit
that referenced
this pull request
Apr 2, 2025
this works with op_mul, which is vectorized-friendly, but doesn't work when we roll out to pattern.h because those ops will not work with Vectorized yet. See TODO in elementwise_util.h ghstack-source-id: 033b63ce3bee8a0136efdab3e03905cafb79b915 ghstack-comment-id: 2738665976 Pull Request resolved: #9432
kimishpatel
reviewed
Apr 2, 2025
| template <typename T> \ | ||
| auto func_name(at::vec::Vectorized<T> vec) { \ | ||
| if constexpr (!::executorch::runtime::is_floating_point<T>::value) { \ | ||
| return at::vec::convert<float>(vec).func_name(); \ |
There was a problem hiding this comment.
is this a valid thing to do? that is convert say an instance of at::vec::Vectorized<int8_t> to float and apply the func_name? Maybe i am misunderstanding how this works
| */ | ||
| #define ET_INTERNAL_VECTORIZED_FLOAT_UNARY_FUNC(func_name) \ | ||
| namespace executorch { \ | ||
| inline namespace math { \ |
There was a problem hiding this comment.
what does inline to namespace do?
Comment on lines
+62
to
+63
| * corresponding operator is a "float op" in TensorIterator parlance | ||
| * (i.e., uses something like build_borrowing_binary_float_op()), |
There was a problem hiding this comment.
I dont know if anyone reading this code would understand what this means?
There was a problem hiding this comment.
But this provides answer to my earlier question
| @@ -47,7 +47,7 @@ Tensor& where_out( | |||
| CTYPE_COMPUTE, | |||
| op_name, | |||
| utils::SupportedTensorDtypes::SAME_AS_COMMON>( | |||
| [](const auto val_a, const auto val_b, const auto val_c) { | |||
| [](const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b, const CTYPE_COMPUTE val_c) { | |||
There was a problem hiding this comment.
is this unrelated change?
There was a problem hiding this comment.
no, we can't vectorize this op
| @@ -51,6 +56,34 @@ inline int64_t scalar_to<int64_t>(const Scalar& s) { | |||
| } | |||
|
|
|||
| namespace internal { | |||
| template <typename Ignore, typename T> | |||
| using ignore_first_yield_second = T; | |||
| if constexpr (std::is_invocable_v< | ||
| Op, | ||
| ignore_first_yield_second<Args, Vec>...>) { | ||
| // For bool, we will get a false positive if we rely on only the |
There was a problem hiding this comment.
you mean bool return type?
There was a problem hiding this comment.
bool ctype_compute
Comment on lines
+103
to
+105
| ((inputs.first->scalar_type() == | ||
| CppTypeToScalarType<CTYPE_COMPUTE>::value) && | ||
| ...)); |
There was a problem hiding this comment.
Man, this is a good reminder of all the template meta programming magic
| ...); | ||
| if (!any_is_broadcasted) { | ||
| using Vec = at::vec::Vectorized<CTYPE_COMPUTE>; | ||
| ::executorch::extension::parallel_for( |
There was a problem hiding this comment.
I think doing this blindly for each op is a bit risky in that, no all multithreading is always better. some ops benefit from smaller grain size while others with larger
There was a problem hiding this comment.
IIRC xnnpack blindly parallelizes absolutely everything; we're doing strictly better here
There was a problem hiding this comment.
yeah I am not comparing with xnnpack. In fact one bit part of the reason why we ended up leveraging optimized op lib for some of the llama stuff for exactly that reason. That it blindly parallelized everything and that actually hurt perf
Comment on lines
+125
to
+126
| const auto vectorized_begin = | ||
| begin + (Vec::size() - begin % Vec::size()) % Vec::size(); |
There was a problem hiding this comment.
Feels like something that has chances of bug. Hope we test this enough. I would doubt if our test cases will exercise this code path.
There was a problem hiding this comment.
although I do see you treat scalar left overs of both head and tails separately
| inputs.first->sizes(), out.sizes()) && | ||
| ...); | ||
| if (!any_is_broadcasted) { | ||
| using Vec = at::vec::Vectorized<CTYPE_COMPUTE>; |
There was a problem hiding this comment.
So the one point of contention for me is that why do we need vectorized_math.h which largely is doing trampoline to underlying vectorized methods. Mainly you dont even need to use can_use_vectorized, because on non accelerated platforms Vectorized falls back to scalar impl even if Vec::size() != `. Maybe you said that the generated code would be worse if forced Vectorized, but I am not sure why. Rest makes sense.
However, place where I can potentially see that being useful is for dtype that Vectorized doesnt support but for float I am not sure. So maybe if you can clarify that it would help.
There was a problem hiding this comment.
why do we need vectorized_math.h which largely is doing trampoline to underlying vectorized methods
without it, you can't take the same lambda you already wrote for scalars and reuse it for Vectorized (the change isn't zero because you have to point at executorch::math, but crucially it doesn't require separate code)
|
Ok we synced offline. So the major value prop as I understood: You can write your lambdas without having to explicitly use Vectorized. Why explicitly using Vectorized might not be good? Because maybe Vectorized does not have everything you need to implement your lambda. So as op author you dont have to worry about Vectorized when writing your lambda (although do have to use executorch::math ops). And later if Vectorized added support for the ::math.. ops you get vectorization for free without having to go back and rewrite your lambda with Vectorized. So this is nice. |
swolchok
added a commit
that referenced
this pull request
Apr 23, 2025
Mixed dtype should be uncommon. Here is how we can specialize for the common case. Prepares us to tackle #9241 . Test Plan: automated tests on this PR verify we didn't break the now-deprecated runtime_out_dtypes mode; tests on the next PR will verify that everything works after migration. Also included migration for exactly one operator, op_mul, to verify that the new code compiles. To check performance, I edited examples/models/toy_model/model.py so that MulModule used inputs of size 3000, 2000 instead of 3, 2. I exported it with `python3 -m examples.portable.scripts.export --model_name mul` and saved the resulting `mul.pte`. Then I built in release mode with optimized kernels on, but with mul.out removed from kernels/optimized/optimized.yaml, so that we would use the optimized_portable_kernels build of kernels/portable/op_mul.cpp. Finally, I ran 3 trials on my M1 Macbook Pro using `cmake-out/executor_runner --model_path mul3kby2k.pte --num_executions 1000 --cpu_threads 2`. Resulting times for 1000 iterations in ms: Previous diff: 8295, 8187, 8139 This diff: 2953, 2806, 2861 (For comparison, the actual optimized mul kernel took around 1000 ms to run 1000 iterations, and #9432 later in the stack arrived at similar numbers.)
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This is a first cut at #9241 . In this PR I've vectorized op_mul to make sure that vectorization doesn't break tests; a follow-up PR will make all existing portable ops vectorized-capable. I've left covering ops that use the
unary_ufunc_*utilities in pattern.h for a follow-up push, because pattern.h and elementwise_util need some work before we can migrate pattern.h's utilities to be backed by elementwise_util.